from typing import Dict, Optional
from langchain_core.messages import SystemMessage, HumanMessage
from config.crypto_config import CODE_GEN_PROMPT_TEMPLATES, SYSTEM_MESSAGE_TEMPLATES

class CodeGeneratorService:
    """代码生成服务类，实现核心的代码生成逻辑"""
    
    def __init__(self):
        self.code_gen_templates = CODE_GEN_PROMPT_TEMPLATES
    
    def generate_crypto_eval_code(self, 
                                key_data: str, 
                                gen_type: str, 
                                crypto_eval_extract: Dict[str, str], 
                                api_info: str, 
                                gen_language: str = "python") -> str:
        """
        生成密码测评代码
        
        Args:
            key_data (str): 关键数据类型，如"KEY"
            gen_type (str): 生成类型，如"ACCESS_CONTROL"
            crypto_eval_extract (Dict[str, str]): 提取的密码测评信息
            api_info (str): API信息
            gen_language (str): 目标编程语言
            
        Returns:
            str: 生成的代码
        """
        try:
            # 获取系统消息
            system_message = SystemMessage(content=SYSTEM_MESSAGE_TEMPLATES)
            
            # 处理AUTH类型的特殊情况
            if gen_type == "AUTH":
                return self.code_gen_templates[f"{gen_type}_CODE_GENERATOR"]
            
            # 构建提示词
            prompt_template = self.code_gen_templates.get(f"{gen_type}_CODE_GENERATOR")
            if not prompt_template:
                raise ValueError(f"未找到类型 {gen_type} 的代码生成模板")
            
            # 格式化提示词
            prompt_message = prompt_template.format(
                key_information=crypto_eval_extract.get(key_data, ""),
                crypto_eval_information=crypto_eval_extract.get(gen_type, ""),
                api_information=api_info,
                language=gen_language
            )
            
            # 构建消息列表
            query_messages = [
                system_message,
                HumanMessage(content=prompt_message)
            ]
            
            # 调用LLM生成代码
            llm = self.llm_service.get_llm()
            response = llm(query_messages)
            
            return response.content
            
        except Exception as e:
            raise Exception(f"代码生成失败: {str(e)}")
    
    def generate_code_with_api_check(self, 
                                   key_data: str, 
                                   gen_type: str, 
                                   crypto_eval_extract: Dict[str, str], 
                                   gen_language: str = "python") -> Optional[str]:
        """
        带API检查的代码生成
        
        Args:
            key_data (str): 关键数据类型
            gen_type (str): 生成类型
            crypto_eval_extract (Dict[str, str]): 提取的密码测评信息
            gen_language (str): 目标编程语言
            
        Returns:
            Optional[str]: 生成的代码，如果没有对应的API信息则返回None
        """
        # 获取API信息
        api_info = self.api_info_service.get_api_info(gen_type)
        
        if api_info:
            return self.generate_crypto_eval_code(
                key_data, gen_type, crypto_eval_extract, api_info, gen_language
            )
        else:
            return None
    
    def get_supported_generation_types(self) -> list:
        """
        获取支持的代码生成类型
        
        Returns:
            list: 支持的生成类型列表
        """
        supported_types = []
        for key in self.code_gen_templates.keys():
            if key.endswith("_CODE_GENERATOR"):
                gen_type = key.replace("_CODE_GENERATOR", "")
                supported_types.append(gen_type)
        return supported_types
    
    def get_supported_languages(self) -> list:
        """
        获取支持的编程语言
        
        Returns:
            list: 支持的编程语言列表
        """
        return ["python", "java", "javascript", "csharp", "go"]
    
    def validate_generation_params(self, 
                                 key_data: str, 
                                 gen_type: str, 
                                 crypto_eval_extract: Dict[str, str], 
                                 gen_language: str) -> bool:
        """
        验证代码生成参数
        
        Args:
            key_data (str): 关键数据类型
            gen_type (str): 生成类型
            crypto_eval_extract (Dict[str, str]): 提取的密码测评信息
            gen_language (str): 目标编程语言
            
        Returns:
            bool: 参数是否有效
        """
        # 检查生成类型是否支持
        supported_types = self.get_supported_generation_types()
        if gen_type not in supported_types:
            return False
        
        # 检查编程语言是否支持
        supported_languages = self.get_supported_languages()
        if gen_language not in supported_languages:
            return False
        
        # 检查必要的数据是否存在
        if not crypto_eval_extract or key_data not in crypto_eval_extract:
            return False
        
        return True 